#!/usr/bin/perl
# This script is a modification of one distributed with the CodingQuarry package by Alison Renwick Testa
# It works on a limited GFF3 file, so may need modification
# Generate FASTA file of the mRNAs and proteins, where the names contain the locations.
#
use strict;
use warnings;
use Bio::Seq;
use Bio::SeqIO;
use Bio::DB::Fasta;

if (@ARGV < 3) {
	print "\nUsage: extractCoding.pl <genome fasta file> <GFF3 file> <output prefix>\n\n";
	exit;
}
print ("Executing extractCoding\n");
die ("Genome file '$ARGV[0]' does not exist ") if (!(-e $ARGV[0])); 
die ("GFF3 file '$ARGV[1]' does not exist ") if (!(-e $ARGV[1])); 

$| = 1;    # Flush output
my $outfile_cds = Bio::SeqIO->new( -format => 'fasta', -file => ">$ARGV[2].cds.fasta" );
my $outfile_pep = Bio::SeqIO->new( -format => 'fasta', -file => ">$ARGV[2].pep.fasta" );

### First, index the genome
my $file_fasta = $ARGV[0];
my $db = Bio::DB::Fasta->new($file_fasta);
print ("Genome fasta parsed\n");

### Second, parse the GFF3
my @mRNA;
my $mRNA_name;
my $frame;
my %root;
my $cnt=2;
my $nGene=1;
open GFF, "<$ARGV[1]" or die $!;

while ( my $line = <GFF> ) {
    chomp $line;
    my @array = split( "\t", $line );
    my $type = $array[2];
    next if ( $type eq 'exon' or $type eq 'UTR' );
	 
    if ( ( $type eq 'gene' ) and ( $. > 2 ) ) { 
        # Collect CDSs and extract sequence of the previous mRNA
		my $mRNA_seq;
        foreach my $coord (@mRNA) {
            my @cds_coord = split( " ", $coord );
            my $cds_seq = $db->seq( $cds_coord[0], $cds_coord[1], $cds_coord[2] );
            $mRNA_seq .= $cds_seq;
        }

        my $output_nucleotide = Bio::Seq->new(
            -seq        => $mRNA_seq,
            -id         => $mRNA_name,
            -display_id => $mRNA_name,
            -alphabet   => 'dna',
	    );
        if ($frame eq '-') {
            $output_nucleotide = $output_nucleotide->revcom();
        }
        my $output_protein = $output_nucleotide->translate();
        $outfile_cds->write_seq($output_nucleotide);
        $outfile_pep->write_seq($output_protein);

        # Now initialize the next mRNA
        if (!exists($root{$array[0]})) {
			$root{$array[0]} = "SC_" . $cnt;
			$cnt++;
		}
        $mRNA_name = $root{$array[0]} . ":" . $array[3] . "-" . $array[4] . "(" . $array[6] . ")";
        $frame=$array[6];
        @mRNA = (); # Empty the mRNA
        $nGene++;
        print "Seqids: " . $cnt . " mRNAs: " . $nGene . "\r" if ($nGene % 1000==0);
    }
    elsif ( $type eq 'gene' ) {    # First mRNA
		$root{$array[0]} = "SC_1";
		$mRNA_name = $root{$array[0]} . ":" . $array[3] . "-" . $array[4] . "(" . $array[6] . ")";
        $frame=$array[6];
    }
    elsif ( $type eq 'CDS' ) {
        my $cds_coord = $array[0] . " " . $array[3] . " " . $array[4];
        push( @mRNA, $cds_coord );
    }
}
print "Seqids: " . $cnt . " mRNAs: " . $nGene . "\n";

#print the last one
my $mRNA_seq;
foreach my $coord (@mRNA) {
    my @cds_coord = split( " ", $coord );
    my $cds_seq = $db->seq( $cds_coord[0], $cds_coord[1], $cds_coord[2] );
    $mRNA_seq .= $cds_seq;
}

my $output_nucleotide = Bio::Seq->new(
            -seq        => $mRNA_seq,
            -id         => $mRNA_name,
            -display_id => $mRNA_name,
            -alphabet   => 'dna',
    );
if ($frame eq '-') {
    $output_nucleotide = $output_nucleotide->revcom();
}
my $output_protein = $output_nucleotide->translate();
$outfile_cds->write_seq($output_nucleotide);
$outfile_pep->write_seq($output_protein);

close GFF;
print "Finish extractCoding\n";

